{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\n# Dynamo Compile Advanced Usage\n\nThis interactive script is intended as an overview of the process by which `torch_tensorrt.dynamo.compile` works, and how it integrates with the new `torch.compile` API.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Imports and Model Definition\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\nfrom torch_tensorrt.dynamo.backend import create_backend\nfrom torch_tensorrt.fx.lower_setting import LowerPrecision"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# We begin by defining a model\nclass Model(torch.nn.Module):\n    def __init__(self) -> None:\n        super().__init__()\n        self.relu = torch.nn.ReLU()\n\n    def forward(self, x: torch.Tensor, y: torch.Tensor):\n        x_out = self.relu(x)\n        y_out = self.relu(y)\n        x_y_out = x_out + y_out\n        return torch.mean(x_y_out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compilation with `torch.compile` Using Default Settings\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Define sample float inputs and initialize model\nsample_inputs = [torch.rand((5, 7)).cuda(), torch.rand((5, 7)).cuda()]\nmodel = Model().eval().cuda()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Next, we compile the model using torch.compile\n# For the default settings, we can simply call torch.compile\n# with the backend \"tensorrt\", and run the model on an\n# input to cause compilation, as so:\noptimized_model = torch.compile(model, backend=\"tensorrt\")\noptimized_model(*sample_inputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compilation with `torch.compile` Using Custom Settings\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Define sample half inputs and initialize model\nsample_inputs_half = [\n    torch.rand((5, 7)).half().cuda(),\n    torch.rand((5, 7)).half().cuda(),\n]\nmodel_half = Model().eval().cuda()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# If we want to customize certain options in the backend,\n# but still use the torch.compile call directly, we can call the\n# convenience/helper function create_backend to create a custom backend\n# which has been pre-populated with certain keys\ncustom_backend = create_backend(\n    lower_precision=LowerPrecision.FP16,\n    debug=True,\n    min_block_size=2,\n    torch_executed_ops={},\n)\n\n# Run the model on an input to cause compilation, as so:\noptimized_model_custom = torch.compile(model_half, backend=custom_backend)\noptimized_model_custom(*sample_inputs_half)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Cleanup\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Finally, we use Torch utilities to clean up the workspace\ntorch._dynamo.reset()\n\nwith torch.no_grad():\n    torch.cuda.empty_cache()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}